# encoding=utf-8
from tqdm import tqdm
import os

import xml.etree.ElementTree as ET


# 新建xml文件
def build_news_xml_file():
    # 设置一个新节点，并设置其标签为root
    root = ET.Element("root")

    # 在root下新建两个子节点,设置其名称分别为sina和chinabyte
    sina = ET.SubElement(root, "sina")
    chinabyte = ET.SubElement(root, "chinabyte")

    # 在sina下新建两个子节点，设置其节点名称分别为number和first
    sina_number = ET.SubElement(sina, "number")
    # sina_number.text = "1"
    sina_first = ET.SubElement(sina, "first")
    sina_first.text = "http://roll.tech.sina.com.cn/internet_all/index_1.shtml"

    # 在sina.number下新建子节点，设置其节点名称分别为sublabel
    sublabel_01 = ET.SubElement(sina_number, "sublabel")
    sublabel_01.text = "sublabel"

    # 在chinabyte下新建两个子节点，设置其节点名称为number和first
    chinabyte_number = ET.SubElement(chinabyte, "number")
    chinabyte_number.text = "1"
    chinabyte_first = ET.SubElement(chinabyte, "first")
    chinabyte_first.text = "http://www.chinabyte.com/more/124566.shtml"

    # 将节点数信息保存在ElementTree中，并且保存为XML格式文件
    tree = ET.ElementTree(root)
    tree.write("urlfile.xml")


def update_tag_content(xml_file, tag_name, tag_content, new_content):
    file_tree = ET.parse(xml_file)
    tree_root = file_tree.getroot()
    for subLabel in tree_root.iter(tag_name):
        # new_content = "guest"
        if subLabel.text == tag_content:
            subLabel.text = new_content
    file_tree.write(xml_file)


def add_object_content(xml_file, cls_box: dict):
    file_tree = ET.parse(xml_file)
    tree_root = file_tree.getroot()
    for name, boxes in cls_box.items():
        for box in boxes:
            object_det = ET.SubElement(tree_root, "object")
            object_name = ET.SubElement(object_det, "name")
            object_name.text = name
            object_pose = ET.SubElement(object_det, "pose")
            object_pose.text = "Unspecified"
            object_truncated = ET.SubElement(object_det, "truncated")
            object_truncated.text = "0"
            object_difficult = ET.SubElement(object_det, "difficult")
            object_difficult.text = "0"
            object_bndbox = ET.SubElement(object_det, "bndbox")
            object_xmin = ET.SubElement(object_bndbox, "xmin")
            object_xmin.text = str(int(box[0]))
            object_ymin = ET.SubElement(object_bndbox, "ymin")
            object_ymin.text = str(int(box[1]))
            object_xmax = ET.SubElement(object_bndbox, "xmax")
            object_xmax.text = str(int(box[2]))
            object_ymax = ET.SubElement(object_bndbox, "ymax")
            object_ymax.text = str(int(box[3]))
    file_tree.write(xml_file)


def delete_classes(xml_file, classes_list):
    file_tree = ET.parse(xml_file)
    tree_root = file_tree.getroot()
    object_element = tree_root.findall("object")
    for obj in object_element:
        if obj[0].text in classes_list:
            tree_root.remove(obj)
    file_tree.write(xml_file)


def xml_combine(xml_file_01, xml_file_02):
    """
    :param xml_file_01:
    :param xml_file_02:
    :return:
    """
    # 解析第一个xml文件
    file_tree_01 = ET.parse(xml_file_01)
    tree_root_01 = file_tree_01.getroot()
    # 解析第二个xml文件
    file_tree_02 = ET.parse(xml_file_02)
    tree_root_02 = file_tree_02.getroot()
    object_element_02 = tree_root_02.findall("object")
    for obj_02 in object_element_02:
        tree_root_01.append(obj_02)
    file_tree_01.write(xml_file_01)


def xml_read(xml_file):
    xml_tree = ET.parse(xml_file)
    print(type(xml_tree))
    xml_root = xml_tree.getroot()
    print(xml_root[0].text)


def change_pic_size(storage_path):
    """
    修改labelImg读取错误的图片大小
    :param storage_path:
    :return:
    """
    xml_list = os.listdir(storage_path)
    for xml in xml_list:
        abs_path = os.path.join(storage_path, xml)
        xml_tree = ET.parse(abs_path)
        xml_root = xml_tree.getroot()
        size = xml_root.find("size")
        width = size.find("width")
        height = size.find("height")
        if width.text == "640" or height.text == "480":
            print(width.text, "   ", height.text)
            print(abs_path)
            # width.text = "1920"
            # height.text = "1080"
        # xml_tree.write(abs_path)


def check_duplicate(xml_file):
    xml_tree = ET.parse(xml_file)
    xml_root = xml_tree.getroot()
    object_storage = []
    object_ele = xml_root.findall("object")
    for obj in object_ele:
        obj_name = obj.find("name").text
        obj_bndbox = obj.find("bndbox")
        obj_xmin = obj_bndbox[0].text
        obj_ymin = obj_bndbox[1].text
        obj_xmax = obj_bndbox[2].text
        obj_ymax = obj_bndbox[3].text
        # obj_list = [obj_name, obj_xmin, obj_ymin, obj_xmax, obj_ymax]
        obj_list = [obj_xmin, obj_ymin, obj_xmax, obj_ymax]
        if obj_list not in object_storage:
            object_storage.append(obj_list)
        else:
            print(xml_file, "\n", obj_list)


def delete_duplicate(xml_file):
    xml_tree = ET.parse(xml_file)
    xml_root = xml_tree.getroot()
    object_storage = []
    object_ele = xml_root.findall("object")
    for obj in object_ele:
        obj_name = obj.find("name").text
        obj_bndbox = obj.find("bndbox")
        obj_xmin = obj_bndbox[0].text
        obj_ymin = obj_bndbox[1].text
        obj_xmax = obj_bndbox[2].text
        obj_ymax = obj_bndbox[3].text
        obj_list = [obj_xmin, obj_ymin, obj_xmax, obj_ymax]
        if obj_list not in object_storage:
            object_storage.append(obj_list)
        else:
            print(xml_file, "\n", obj_list)
            xml_root.remove(obj)
    xml_tree.write(xml_file)


if __name__ == "__main__":
    storage_path = r"/home/kaijia/algo-env/Utils/sample/217A/Annotations"
    # change_pic_size(storage_path)
    xml_list = os.listdir(storage_path)
    # for xml in tqdm(xml_list):
    #     # print("Processing ", xml)
    #     abs_path = os.path.join(storage_path, xml)
    #     check_duplicate(abs_path)
    #     delete_duplicate(abs_path)
    # #合并xml文件
    # xml_directory_01 = r"/home/kaijia/algo-env/datasets/hq_safeguard_dataset/Annotations-01"
    # xml_directory_02 = r"/home/kaijia/algo-env/datasets/hq_safeguard_dataset/Annotations-02"
    # xml_list_01 = os.listdir(xml_directory_01)
    # for xml_file in xml_list_01:
    #     xml_file_01 = os.path.join(xml_directory_01, xml_file)
    #     xml_file_02 = os.path.join(xml_directory_02, xml_file)
    #     xml_combine(xml_file_01, xml_file_02)

    ########## 删除对应类别 ##########
    # storage_path = "/home/stark/algo-env/datasets/airport-camera/0008/ReID/Annotations"
    # xml_list = os.listdir(storage_path)
    # classes_list = ["tire"]
    # for xml in xml_list:
    #     xml_file = os.path.join(storage_path, xml)
    #     delete_classes(xml_file, classes_list)
    ########## 修改对应类别 ##########
    # xml_list = os.listdir(storage_path)
    # for xml in xml_list:
    #     xml_file = os.path.join(storage_path, xml)
    #     modify_tag_content(xml_file, "name", "person", "guest")
    ########## 增加bndbox标签内容 ##########
    for xml in tqdm(xml_list):
        abs_path = os.path.join(storage_path, xml)
        cls_box = {"tire": [[975, 298, 1010, 331]]}
        add_object_content(abs_path, cls_box)
